from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from transformers import AutoTokenizer,AutoModelForCausalLM,pipeline
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.callbacks import StreamingStdOutCallbackHandler
import os
from langchain_core.callbacks import StreamingStdOutCallbackHandler,CallbackManager
#创建一个聊天模型
class ChatModel:
    #定义初始化函数
    def __init__(self):
        #获取环境变量
        load_dotenv()
    #获取在线模型
    def get_line_model(self):
        model=ChatOpenAI(model=os.getenv("model_name"),temperature=0.7)
        return model

    #获取本地模型
    def get_local_model(self):
        #加载模型
        model_path=os.getenv("base_model")
        #加载分词器
        tokenizer=AutoTokenizer.from_pretrained(model_path)
        #加载权重
        model=AutoModelForCausalLM.from_pretrained(model_path)
        #创建管道
        pipe=pipeline("text-generation",model=model,tokenizer=tokenizer,max_length=512,temperature=0.7)
        return pipe

    #获取嵌入模型
    def get_embedding_model(self):
        #加载模型
        model=DashScopeEmbeddings(
            model=os.getenv("embedding_model"),
            dashscope_api_key=os.getenv("dash_scope_key")
        )
        return model

    #获取流式模型
    def get_line_stream_model(self):
        callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
        model=ChatOpenAI(model=os.getenv("model_name"),
                         temperature=0.7,
                         streaming=True,# 启用流式传输
                         callbacks=callback_manager,# 流式输出到控制台
                         verbose=True
                         )
        return model